import pandas as pd
import numpy as np
import lightgbm as lgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, roc_curve, classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import warnings

# --- 配置 ---
# 忽略一些未来版本的警告，使输出更整洁
warnings.filterwarnings('ignore', category=FutureWarning)
# 设置pandas显示选项
pd.set_option('display.max_columns', None)

# --- 1. 数据加载 ---
try:
    df = pd.read_csv('../data/train.csv')
    print(f"成功加载 train.csv 文件，数据包含 {df.shape[0]} 行和 {df.shape[1]} 列。")
except FileNotFoundError:
    print("错误：找不到 'train.csv' 文件。请确保文件与脚本在同一目录下。")
    exit()


# --- 2. 特征工程 ---
def feature_engineer(data):
    """在一个DataFrame上执行特征工程"""
    df_fe = data.copy()

    # 防止除以零，添加一个极小值
    epsilon = 1e-6

    # a. 比例特征
    df_fe['IncomePerYearOfExperience'] = df_fe['MonthlyIncome'] / (df_fe['TotalWorkingYears'] + epsilon)
    df_fe['YearsAtCompanyRatio'] = df_fe['YearsAtCompany'] / (df_fe['TotalWorkingYears'] + epsilon)
    df_fe['PromotionToWorkingYearRatio'] = df_fe['YearsSinceLastPromotion'] / (df_fe['TotalWorkingYears'] + epsilon)

    # b. 聚合特征
    satisfaction_cols = ['EnvironmentSatisfaction', 'JobSatisfaction', 'RelationshipSatisfaction', 'WorkLifeBalance']
    df_fe['OverallSatisfaction'] = df_fe[satisfaction_cols].sum(axis=1)

    # c. 交互/派生特征将在数据划分后进行，以避免数据泄露

    return df_fe


print("\n开始进行特征工程...")
df_engineered = feature_engineer(df)
print(
    "特征工程完成。新增特征：'IncomePerYearOfExperience', 'YearsAtCompanyRatio', 'PromotionToWorkingYearRatio', 'OverallSatisfaction'")

# --- 3. 数据预处理与划分 ---
# 定义特征(X)和目标(y)
y = df_engineered['Attrition']
X = df_engineered.drop('Attrition', axis=1)

# 丢弃无用列
cols_to_drop = ['EmployeeNumber', 'StandardHours', 'Over18']
X = X.drop(cols_to_drop, axis=1)

# 划分训练集和测试集 (80% 训练, 20% 测试)
# stratify=y 确保训练集和测试集中流失员工的比例与原始数据一致
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"\n数据已划分为训练集 ({X_train.shape[0]} 条) 和测试集 ({X_test.shape[0]} 条)。")

# --- 4. 避免数据泄露的特征工程 (在数据划分后进行) ---
# 计算训练集中每个JobRole的平均收入
role_income_map = X_train.groupby('JobRole')['MonthlyIncome'].mean().to_dict()

# 应用到训练集和测试集
X_train['IncomeToRoleAverageRatio'] = X_train.apply(
    lambda row: row['MonthlyIncome'] / role_income_map.get(row['JobRole'], row['MonthlyIncome']), axis=1
)
X_test['IncomeToRoleAverageRatio'] = X_test.apply(
    lambda row: row['MonthlyIncome'] / role_income_map.get(row['JobRole'], row['MonthlyIncome']), axis=1
)
print("新增特征 'IncomeToRoleAverageRatio'，已在训练集和测试集上应用。")

# --- 5. 最终预处理 (转换数据类型) ---
categorical_features = X_train.select_dtypes(include=['object']).columns
for col in categorical_features:
    X_train[col] = X_train[col].astype('category')
    # 确保测试集的类别与训练集一致
    X_test[col] = pd.Categorical(X_test[col], categories=X_train[col].cat.categories)

# --- 6. 模型训练 ---
print("\n开始训练LightGBM模型...")
lgbm = lgb.LGBMClassifier(objective='binary', random_state=42, n_estimators=200, learning_rate=0.05, num_leaves=20)
lgbm.fit(X_train, y_train)
print("模型训练完成。")

# --- 7. 模型评估 ---
# 在测试集上预测概率
y_pred_proba = lgbm.predict_proba(X_test)[:, 1]
y_pred = lgbm.predict(X_test)

# 计算AUC
auc_score = roc_auc_score(y_test, y_pred_proba)

print("\n--- 模型性能评估 (在测试集上) ---")
print(f"AUC Score: {auc_score:.4f}")

if auc_score >= 0.75:
    print("目标达成：AUC值 >= 0.75！")
else:
    print("目标未达成：AUC值 < 0.75。可以尝试调整模型参数或进行更深入的特征工程。")

# 打印分类报告
print("\n分类报告 (Classification Report):")
print(classification_report(y_test, y_pred, target_names=['Not Attrition', 'Attrition']))

# 绘制ROC曲线
fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)
plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (AUC = {auc_score:.4f})')
plt.plot([0, 1], [0, 1], color='gray', lw=2, linestyle='--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('Receiver Operating Characteristic (ROC) Curve')
plt.legend(loc="lower right")
plt.grid(True)
# plt.show()
plt.savefig('../data/roc_curve.png')
# --- 8. 特征重要性分析 ---
feature_importance_df = pd.DataFrame({
    'feature': X_train.columns,
    'importance': lgbm.feature_importances_
}).sort_values('importance', ascending=False)

plt.figure(figsize=(12, 10))
sns.barplot(x='importance', y='feature', data=feature_importance_df.head(20))
plt.title('Top 20 Feature Importances')
plt.tight_layout()
# plt.show()
plt.savefig('../data/feature_importance.png')

print("\n--- 特征重要性分析 ---")
print("Top 10 最重要的特征:")
print(feature_importance_df.head(10))
